use std::fs;
use std::env;
use std::process::Command;
use syn::{File, Item, ItemFn};
use quote::ToTokens; // 导入 ToTokens trait
use anyhow::{Result, Context};

fn main() -> Result<()> {
    // 获取命令行参数
    let args: Vec<String> = env::args().collect();
    if args.len() != 3 {
        eprintln!("Usage: {} <temp_file> <c_function_name>", args[0]);
        std::process::exit(1);
    }

    let temp_file = &args[1];
    let c_function_name = &args[2];

    // 读取临时文件中的 Rust 代码
    let code = fs::read_to_string(temp_file)
        .with_context(|| format!("Failed to read file: {}", temp_file))?;

    // 使用 syn 解析代码
    let ast = match syn::parse_str::<File>(&code) {
        Ok(ast) => ast,
        Err(_) => {
            // 如果解析失败，输出空
            println!("");
            return Ok(());
        }
    };

    // 提取目标函数
    let mut target_function = None;
    for item in ast.items {
        if let Item::Fn(item_fn) = item {
            if item_fn.sig.ident == c_function_name {
                target_function = Some(item_fn);
                break;
            }
        }
    }

    // 如果没有找到目标函数，输出空
    if target_function.is_none() {
        println!("");
        return Ok(());
    }

    // 将目标函数转换为字符串
    let target_function_code = item_fn_to_string(&target_function.unwrap());

    // 使用 rustfmt 格式化代码
    let formatted_code = format_with_rustfmt(&target_function_code)?;

    // 输出格式化后的代码
    println!("{}", formatted_code);

    Ok(())
}

/// 将 `ItemFn` 转换为字符串
fn item_fn_to_string(item_fn: &ItemFn) -> String {
    let mut code = String::new();

    // 添加函数属性（如果有）
    for attr in &item_fn.attrs {
        code.push_str(&attr.to_token_stream().to_string());
        code.push('\n'); // 每个属性单独一行
    }

    // 添加函数签名和函数体
    code.push_str(&item_fn.sig.to_token_stream().to_string()); // 包含函数名
    code.push_str(&item_fn.block.to_token_stream().to_string());

    code
}

/// 使用 rustfmt 格式化代码
fn format_with_rustfmt(code: &str) -> Result<String> {
    let mut child = Command::new("rustfmt")
        .arg("--emit=stdout")
        .arg("--edition=2021")
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .spawn()
        .with_context(|| "Failed to spawn rustfmt process")?;

    // 将代码写入 rustfmt 的标准输入
    if let Some(mut stdin) = child.stdin.take() {
        use std::io::Write;
        stdin.write_all(code.as_bytes())?;
    }

    // 读取 rustfmt 的标准输出
    let output = child.wait_with_output()?;
    if !output.status.success() {
        anyhow::bail!("rustfmt failed with status: {}", output.status);
    }

    let formatted_code = String::from_utf8(output.stdout)
        .with_context(|| "rustfmt output is not valid UTF-8")?;

    Ok(formatted_code)
}
